"""
LLM健康检查模块

提供了检查LLM连接和功能状态的工具函数
"""

import asyncio
import logging
import time
from typing import Any, Callable, Dict, Optional

logger = logging.getLogger("wren-ai-service")


class LLMHealthChecker:
    """LLM健康检查器，提供了检查LLM状态的方法和定期执行健康检查的功能"""

    def __init__(
        self,
        llm_generator: Callable,
        model_name: str = "unknown",
        api_base: str = "unknown",
        check_interval: int = 300,
    ):
        """
        初始化健康检查器

        Args:
            llm_generator: LLM生成器函数
            model_name: 模型名称
            api_base: API基础URL
            check_interval: 健康检查间隔（秒）
        """
        self.llm_generator = llm_generator
        self.model_name = model_name
        self.api_base = api_base
        self.check_interval = check_interval
        self.task = None
        self.last_check_result = None

    async def check_health(self) -> Dict[str, Any]:
        """
        执行LLM健康检查

        Returns:
            Dict: 包含健康状态的字典
        """
        start_time = time.time()
        result = {
            "status": "unknown",
            "message": "",
            "llm_provider": "litellm_llm",
            "model": self.model_name,
            "api_base": self.api_base,
            "latency_ms": 0,
            "timestamp": start_time,
        }

        try:
            # 执行简单推理任务检查LLM是否响应
            prompt = "Return exactly the 5 letters: HELLO"
            test_response = await self.llm_generator(prompt=prompt)

            # 检查响应是否符合预期
            response_text = (
                test_response["replies"][0]
                if test_response
                and "replies" in test_response
                and test_response["replies"]
                else ""
            )
            if "HELLO" in response_text:
                result["status"] = "ok"
                result["message"] = "LLM响应正常"
            else:
                result["status"] = "degraded"
                result["message"] = f"LLM响应异常: {response_text[:100]}"

            # 添加完整的响应内容作为调试信息
            result["debug"] = {
                "response": response_text,
                "meta": test_response.get("meta", {}),
            }
        except Exception as e:
            logger.exception("LLM健康检查失败")
            result["status"] = "error"
            result["message"] = f"LLM连接错误: {str(e)}"
            result["debug"] = {"error": str(e)}

        # 计算延迟
        result["latency_ms"] = round((time.time() - start_time) * 1000, 2)
        self.last_check_result = result
        return result

    async def _periodic_check_task(self):
        """定期健康检查任务"""
        # 首次启动等待一段时间，确保服务完全初始化
        await asyncio.sleep(30)

        while True:
            try:
                # 执行健康检查
                health_result = await self.check_health()
                status = health_result.get("status", "unknown")
                message = health_result.get("message", "")
                latency = health_result.get("latency_ms", 0)

                # 根据状态记录不同级别的日志
                if status == "ok":
                    logger.info(f"LLM健康状态: {status}, 延迟: {latency}ms, {message}")
                elif status == "degraded":
                    logger.warning(
                        f"LLM健康状态降级: {status}, 延迟: {latency}ms, {message}"
                    )
                else:
                    logger.error(f"LLM健康状态异常: {status}, {message}")

                # 等待指定的时间间隔
                await asyncio.sleep(self.check_interval)
            except asyncio.CancelledError:
                logger.info("LLM健康检查任务已取消")
                break
            except Exception as e:
                logger.exception(f"执行定期LLM健康检查时出错: {str(e)}")
                await asyncio.sleep(60)  # 出错后等待1分钟再尝试

    def start_periodic_check(self):
        """启动定期健康检查任务"""
        if self.task is None or self.task.done():
            self.task = asyncio.create_task(self._periodic_check_task())
            logger.info(f"启动了LLM定期健康检查任务，检查间隔为{self.check_interval}秒")
        return self.task

    def stop_periodic_check(self):
        """停止定期健康检查任务"""
        if self.task and not self.task.done():
            self.task.cancel()
            logger.info("LLM健康检查任务已停止")

    def get_last_check_result(self) -> Optional[Dict[str, Any]]:
        """获取最近一次健康检查结果"""
        return self.last_check_result
